#ifndef SRC_MPI_CIPHER_CANNON
#define SRC_MPI_CIPHER_CANNON

#include "mpi/cipher_mpi_base.hpp"

template <typename Scheme, typename ExecOutput>
class cipher_cannon : public cipher_mpi_base<Scheme, ExecOutput> {
protected:
  /* @brief     Illustrates `AXIS` and `DIRECTIONS`
   * @details   Consider dims[0] = X and dims[1] = Y
   *            D: Down; U: Up; L: Left; R: Right
   * +----> Y
   * |                 LEFT
   * |         <--------------------
   * v       +-----------+-----------+
   * X       |           |           |
   *       ^ | process 0 | process 1 |
   *       | |           |           |
   *    UP | +-----------+-----------+
   *       | |           |           |
   *       | | process 2 | process 3 |
   *         |           |           |
   *         +-----------------------+
   **/
  enum AXIS { X, Y, AXIS_NUM };
  enum DIRECTIONS { D, U, L, R, DIRECT_NUM };
  struct metric : public cipher_mpi_base<Scheme, ExecOutput>::metric {
    // * exec_time  <- cipher_mpi_base<ExecOutput>::metric
    //   * shift_align_time
    //   * shift_compute_time
    //     * shift step-by-step (sc_comm_time)
    //     * serialization
    //     * local cgemm
    double shift_align_time = 0.0;
    double shift_compute_time = 0.0;
    double sc_comm_time = 0.0;
    double sc_serialize_time = 0.0;
    double sc_cgemm_time = 0.0;
    std::vector<double> cgemm_times;
  };

private:
  cipher_cannon() = delete;

public:
  /**
   * @brief Construct a new `cipher_cannon` object
   *
   * @param root  aka. cipher_root(possess initial data and public key)
   *              which is a rank in `comm`
   * @param comm  Construct `cipher_cannon` in `coom`
   * @param m     dim-m
   * @param n     dim-n
   * @param k     dim-k
   */
  cipher_cannon(int root, MPI_Comm comm, long m, long n, long k);
  virtual ~cipher_cannon() = default;

  long block_m() const { return mb_; }
  long block_n() const { return nb_; }
  long block_k() const { return kb_; }

  void report(std::ostream &out) const override;

protected:
  // MPI overall
  int np_ = 0;
  int rp_ = 0;

  // Problem size
  long m_ = 0;
  long n_ = 0;
  long k_ = 0;
  long mb_ = 0;
  long nb_ = 0;
  long kb_ = 0;

  metric stat_;

  // Cannon's metadata
  MPI_Comm cannon_comm_;
  int sqrt_np_ = 0;
  int sa_rank_[DIRECT_NUM]; // ShiftAlign
  int sc_rank_[DIRECT_NUM]; // ShiftCompute
};

//* template specialization
// for single-ciphertext version
template class cipher_cannon<helib::BGV, std::string>;
// for multi-ciphertext version
template class cipher_cannon<helib::BGV, std::vector<std::string>>;

// * shhmm_cannon
//   * s      : single-ciphertext
//   * h      : hypercube packing
//   * hmm    : homomorphic matrix multiplication
//   * cannon : Cannon's PGEMM algorithm
// * only BGV supports hypercube packing
class shhmm_cannon : public cipher_cannon<helib::BGV, std::string> {
private:
  template <typename T>
  void init_data(int root, int rank, MPI_Comm comm, const T *A, const T *B);

public:
  //* BGV parameters: params<helib::BGV>
  shhmm_cannon(int root, MPI_Comm comm, long m, long n, long k,
               const params<helib::BGV> &params,
               const hmm_status<shmm_engine> &status);

  void setup_crypto(int root, int rank, const params<helib::BGV> &params,
                    const cube_signature &cube, MPI_Comm comm);

  void init_data(const NTL::mat_ZZ &A, const NTL::mat_ZZ &B);
  void init_data(const int *A, const int *B) override;

  void collect_data(const std::string &partial_C,
                    std::vector<NTL::mat_ZZ> &result,
                    bool local = false) override;

  void exec(std::string &partial_C, bool fhe4d = false) override;

  //* hypercube HMM engine: hmm_status<hypercube_hmme>
  hmm_status<hypercube_hmme> *hmme_status() override {
    return static_cast<hypercube_hmme *>(hmm_engine_)->status();
  }
};

// * mhhmm_cannon
//   * m      : multi-ciphertext
//   * h      : hypercube packing
//   * hmm    : homomorphic matrix multiplication
//   * cannon : Cannon's PGEMM algorithm
// * only BGV supports hypercube packing
class mhhmm_cannon
    : public cipher_cannon<helib::BGV, std::vector<std::string>> {
private:
  typedef std::vector<std::vector<helib::Ctxt>> mat_ctxt;

  template <typename T>
  void init_data(int root, int rank, MPI_Comm comm, const T *A, const T *B);

  void strassen(std::vector<helib::Ctxt> &C, const mat_ctxt &A,
                const mat_ctxt &B);
  void stdblk(std::vector<helib::Ctxt> &C, const mat_ctxt &A,
              const mat_ctxt &B);

public:
  //* BGV parameters: params<helib::BGV>
  mhhmm_cannon(int root, MPI_Comm comm, long m, long n, long k, long opt,
               const params<helib::BGV> &params,
               const hmm_status<shmm_engine> &status);

  void init_data(const int *A, const int *B) override;

  void collect_data(const std::vector<std::string> &partial_C,
                    std::vector<NTL::mat_ZZ> &result,
                    bool local = false) override;

  void strassen_state(bool set) { use_strassen_ = set; }
  void exec(std::vector<std::string> &partial_C, bool fhe4d = false) override;

  //* hypercube HMM engine: hmm_status<hypercube_hmme>
  hmm_status<hypercube_hmme> *hmme_status() override {
    return static_cast<hypercube_hmme *>(hmm_engine_)->status();
  }

private:
  bool use_strassen_ = false;

  long mbb_ = 0;
  long nbb_ = 0;
  long kbb_ = 0;

  long m_nblk_ = 0;
  long n_nblk_ = 0;
  long k_nblk_ = 0;
};

#endif /* SRC_MPI_CIPHER_CANNON */
